import sqlite3
from contextlib import closing
from utils.configFinRAG import sql_examples_path
from utils.instances import TOKENIZER, LLM

# 假设这里已经有一个配置好的数据库连接池
# db_pool = ...

def load_sql_examples():
    """加载SQL示例数据"""
    examples = pd.read_csv(sql_examples_path, delimiter=",", header=0)
    questions, sqls, infos, fas, tokens = [], [], [], [], []
    for _, row in examples.iterrows():
        questions.append(row['问题'])
        sqls.append(row['SQL'])
        infos.append(row['资料'])
        fas.append(row['FA'])
        tokens.append(TOKENIZER(row['问题'])['input_ids'])
    return questions, sqls, infos, fas, tokens

def query_database(sql_query):
    """查询数据库并返回结果"""
    with closing(sqlite3.connect('your_database_path')) as conn:
        with conn.cursor() as cs:
            success_flag, result = query_db(sql_query, cs)
    return success_flag, result

def sql_retrieve_chain(query, example_data):
    """主流程函数"""
    result_prompt, sql_query = generate_sql(query, LLM, *example_data)
    success_flag, query_result = query_database(sql_query)
    answer = generate_answer(query, query_result, LLM, *example_data)
    return answer

# 在应用启动时加载一次即可
example_data = load_sql_examples()

# 调用时
answer = sql_retrieve_chain("你的查询语句", example_data)